Security: Unsafe torch.load with weights_only=False enables arbitrary code execution#2645
Conversation
In `CheckpointManager.load_from_path()`, `torch.load()` is called with `weights_only=False` (line with `torch.load(f, weights_only=False)`). This is a critical security vulnerability because it allows arbitrary Python code execution during checkpoint deserialization. An attacker who can control the checkpoint file can execute arbitrary code on the system. The `weights_only=False` parameter disables PyTorch's safe unpickling mode, which was the default behavior in older PyTorch versions but is known to be dangerous. Signed-off-by: tomaioo <203048277+tomaioo@users.noreply.github.com>
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit f322a5f. Configure here.
| else: | ||
| with open(ckpt_path / "progress.pt", "rb") as f: | ||
| state = torch.load(f, weights_only=False) | ||
| state = torch.load(f, weights_only=True) |
There was a problem hiding this comment.
Safe load breaks progress resume
High Severity
load_from_path now calls torch.load with weights_only=True, but save_to_path still writes progress.pt via torch.save containing a live Progress dataclass instance. Restricted unpickling does not allow that custom type unless it is explicitly allowlisted, so resuming orchestrator checkpoints (when progress loading is not skipped) is likely to fail at load time.
Reviewed by Cursor Bugbot for commit f322a5f. Configure here.
1 participant
Summary
Security: Unsafe torch.load with weights_only=False enables arbitrary code execution
Problem
Severity:
Critical| File:src/prime_rl/orchestrator/ckpt.py:L56In
CheckpointManager.load_from_path(),torch.load()is called withweights_only=False(line withtorch.load(f, weights_only=False)). This is a critical security vulnerability because it allows arbitrary Python code execution during checkpoint deserialization. An attacker who can control the checkpoint file can execute arbitrary code on the system. Theweights_only=Falseparameter disables PyTorch's safe unpickling mode, which was the default behavior in older PyTorch versions but is known to be dangerous.Solution
Change
torch.load(f, weights_only=False)totorch.load(f, weights_only=True). If the checkpoint requires custom objects, implement a safe deserialization mechanism with explicit allowlists or use a different serialization format like safetensors.Changes
src/prime_rl/orchestrator/ckpt.py(modified)Note
Medium Risk
Reduces arbitrary-code-execution risk on untrusted checkpoints; resume may fail if legacy progress.pt files relied on unsafe pickle types.
Overview
Hardens orchestrator resume by loading
progress.ptwithtorch.load(..., weights_only=True)instead ofweights_only=False, so checkpoint deserialization no longer runs arbitrary pickle code if a checkpoint file is tampered with.Only
CheckpointManager.load_from_pathinckpt.pychanges; save path and buffer loading are unchanged.Reviewed by Cursor Bugbot for commit f322a5f. Bugbot is set up for automated code reviews on this repo. Configure here.